import os
import numpy as np
from matplotlib import pyplot as plt
from mindspore import load_checkpoint, load_param_into_net, Tensor, Model
import mindspore.dataset.transforms.c_transforms as C
import mindspore.dataset.vision.c_transforms as CV
from mindspore.dataset.vision import Inter
from mindspore import dtype as mstype
import mindspore.dataset as ds

from lenet import LeNet5

data_path = "./datasets/"
raw_data = ds.MnistDataset(os.path.join(data_path, "test"), num_samples=1, shuffle=False)

def draw_data(raw_data):
    for item in raw_data.create_dict_iterator(num_epochs=1, output_numpy=True):
        image = item["image"]
        label = item["label"]
    plt.imshow(image, cmap=plt.cm.gray)
    plt.title(label)
    plt.show()

def transform_data(data, batch_size):
    # 定义所需要操作的map映射
    resize_op = CV.Resize((32, 32), interpolation=Inter.LINEAR)     # 目标将图片大小调整为32*32，这样特征图能保证28*28，和原图一致
    rescale_nml_op = CV.Rescale(1 / 0.3081 , -1 * 0.1307 / 0.3081)  # 数据集的标准化系数
    rescale_op = CV.Rescale(1.0 / 255.0, 0.0)                       # 数据做标准化处理，所得到的数值分布满足正态分布
    hwc2chw_op = CV.HWC2CHW()                                       # 转置操作
    type_cast_op = C.TypeCast(mstype.int32)

    # 使用map映射函数，将数据操作应用到数据集
    data = data.map(operations=type_cast_op, input_columns="label")
    data = data.map(operations=[resize_op, rescale_op, rescale_nml_op, hwc2chw_op], input_columns="image")
     # 进行batch操作
    return data.batch(batch_size)

param_dict = load_checkpoint("checkpoint_lenet-1_1875.ckpt")
# # 加载参数到网络中
net = LeNet5()
load_param_into_net(net, param_dict)
model = Model(net)
# 定义测试数据集，batch_size设置为1，则取出一张图片
for item in transform_data(raw_data, 1):
    image = item[0]
    label = item[1]
# # 使用函数model.predict预测image对应分类
output = model.predict(Tensor(image))
predicted = np.argmax(output.asnumpy(), axis=1)
# 输出预测分类与实际分类
print(f'Predicted: "{predicted[0]}", Actual: "{label[0]}"')
draw_data(raw_data)
